#include <torch/torch.h>
#include <vector>
#include <cstdio>

at::Tensor gdw_conv_cuda_forward(
    at::Tensor x,
    at::Tensor w,
    at::Tensor b);

std::vector<at::Tensor> gdw_conv_cuda_backward(
    at::Tensor grad_y,
    at::Tensor x,
    at::Tensor w,
    at::Tensor b);

#define CHECK_CUDA(x) AT_ASSERT(x.type().is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) AT_ASSERT(x.is_contiguous(), #x " must be contiguous")
#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)


at::Tensor gdw_conv_forward(
    at::Tensor x,
    at::Tensor w,
    at::Tensor b) {
    CHECK_INPUT(x);
    CHECK_INPUT(w);
    CHECK_INPUT(b);
    return gdw_conv_cuda_forward(x, w, b);
}

std::vector<at::Tensor> gdw_conv_backward(
    at::Tensor grad_y,
    at::Tensor x,
    at::Tensor w,
    at::Tensor b) {
    CHECK_INPUT(grad_y);
    CHECK_INPUT(x);
    CHECK_INPUT(w);
    CHECK_INPUT(b);

    return gdw_conv_cuda_backward(grad_y, x, w, b);
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("forward", &gdw_conv_forward, "global depthwise conv forward (CUDA)");
  m.def("backward", &gdw_conv_backward, "global depthwise conv backward (CUDA)");
}